import argparse
import json


def get_subject_mapping():
    mapping_path = "ceval_code/subject_mapping.json"
    with open(mapping_path, 'r') as file:
        mapping = json.load(file)
    return mapping


categories = {'STEM': { 'correct': 0, 'total': 0, 'avg': 0. },
             'Social Science': { 'correct': 0, 'total': 0, 'avg': 0.  },
             'Humanities': { 'correct': 0, 'total': 0, 'avg': 0.  },
             'Other': { 'correct': 0, 'total': 0, 'avg': 0.  }}

def get_subject_cnt():
    with open("subjects_cnt.txt", 'r') as file:
        subjects_cnt = file.readlines()
    subjects_cnt = [line.split() for line in subjects_cnt]
    subjects_cnt = {line[1]: int(line[0]) - 1 for line in subjects_cnt}
    return subjects_cnt


def compute(subject_data, mapping):
    for data in subject_data:
        subject_name = data['subject']
        category = mapping[subject_name][2]
        # ratio = data['ratio']
        correct = data['correct']
        total = data['total']
        categories[category]['correct'] += correct #subject_cnt[subject_name]
        categories[category]['total'] += total #ratio * subject_cnt[subject_name]
        # categories[category]['cnt'] += 1
        # categories[category]['sum'] += ratio

    for key, val in categories.items():
        val['avg'] = val['correct'] / val['total']

    correct = 0
    total = 0
    for key, val in categories.items():
        correct += val['correct']
        total += val['total']
        print(f"{key}: {val['correct']}/{val['total']} {val['avg'] * 100:.1f}%")

    print(f"Avg: {correct}/{total} {correct / total * 100:.1f}%")


def get_subject_data(path):
    correct_result_path = path #"12k_step_ceval.txt"
    with open(correct_result_path, 'r') as file:
        lines = file.readlines()
    data = []
    for line in lines:
        subject_data = line.strip().split()
        data.append({
            'subject': subject_data[-3][:-1],
            'correct': int(subject_data[-2].split('/')[0]),
            'total': int(subject_data[-2].split('/')[1]),
            'ratio': float(subject_data[-1][:-1])
        })
    return data


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', type=str, required=True)
    args = parser.parse_args()
    subject_mapping = get_subject_mapping()
    subject_data = get_subject_data(args.file)
    # print(subjects_cnt)
    # subject_cnt = get_subject_cnt()
    compute(subject_data, subject_mapping)


if __name__ == "__main__":
    main()
